# utf-8
import math
import os
import sys
import configparser
from time import sleep
import paho.mqtt.client as mqtt
import traceback

client = None

lily_cmds = []
lily_docs = {}
lily_hijackers = []
lily_exts = []


def on_connect(client, userdata, flags, rc):
    if rc:
        exit(-rc)


def on_disconnect(client, userdata, flags):
    raise Exception('connection lost')


def mqtt_blob(blob: bytes):
    global blob_id
    blob_id += 1
    n = len(blob)
    batch_size = 32768 - 256
    seg_n = math.ceil(n/batch_size) - 1
    top_ = f'lily/{client_id}/bin/{blob_id}/{seg_n}/'
    f = 0
    t = f + batch_size
    for i in range(seg_n+1):
        top = top_ + str(i)
        info = client.publish(top, blob[f:t])
        sleep(0.1)
        if info.rc != mqtt.MQTT_ERR_SUCCESS:
            print(f'error, {i}:send failed')
        f = t
        if f >= n:
            break
        t = f + batch_size
        if t >= n:
            t = n


def mqtt_blob_str(blob: bytes):
    n = len(blob)
    batch_size = 32768 - 256
    seg_n = math.ceil(n/batch_size)
    top_ = f'lily/{client_id}/msg'
    f = 0
    t = f + batch_size
    if seg_n <= 1:
        client.publish(top_, blob)
        return
    for i in range(seg_n):
        info = client.publish(top_, blob[f:t])
        sleep(0.1)
        if info.rc != mqtt.MQTT_ERR_SUCCESS:
            print(f'error, {i}:send failed')
        f = t
        if f >= n:
            break
        t = f + batch_size
        if t >= n:
            t = n


def mqtt_out(msg: str, end='\n'):
    client.publish(f'lily/{client_id}/msg', msg+end)


def lily_add_hijack(callback):
    lily_hijackers.append(callback)


def lily_add_hijack_ext(callback):
    lily_exts.append(callback)


def lily_do(cmd: str):
    if lily_hijackers:
        go_on = False
        try:
            free, go_on = lily_hijackers[-1](cmd)
            if free:
                lily_hijackers.pop()
        except Exception as e:
            e = str(e)
            mqtt_out(f'error:{e}')
        if not go_on:
            return

    raw_cmd = cmd.strip()
    cmd = raw_cmd.split(' ')
    n = len(cmd)
    if not cmd:
        return
    s = cmd[0]
    for c, f in lily_cmds:
        if c != s:
            continue
        try:
            rtc = f(n, cmd)
            if type(rtc) == int and rtc < 0:
                mqtt_out(f'error return {rtc}')
        except Exception as e:
            e = str(e)
            traceback.print_exc()
            mqtt_out(f'error:{e}')
        return
    mqtt_out(f'cmd "{cmd[0]}" not found')
    return


def lily_do_ext(cmd: str):
    if not lily_exts  and len(lily_exts)<1:
        return
    f = lily_exts[-1]
    try:
        free = f(cmd)
        if free:
            lily_exts.pop()
    except Exception as e:
        e = str(e)
        traceback.print_exc()
        mqtt_out(f'error:{e}')


def new_cmd(name: str, f, doc=''):
    lily_cmds.append((name, f))
    if doc:
        lily_docs[name] = doc


def on_message(client_name, userdata, msg):
    topic = msg.topic
    cmd = msg.payload.decode('utf-8', 'ignore')
    if topic == mqtt_subscribe_topics[0]: # cmd
        lily_do(cmd)
    elif topic == mqtt_subscribe_topics[1]: # ext
        lily_do_ext(cmd)
    elif topic == mqtt_subscribe_topics[2]: # ping
        client.publish(f'lily/{client_id}/pong', cmd)

pty_doc = '''pty [-r] [-e] [-t term] [row col]: start a pty
    -h: help
    -r: raw input, no newline
    -e: local echo
    -t: terminal type, default xterm-256color
    row col: set terminal size
'''


def cmd_pty(n, args):
    import jc_pty
    raw_input = False
    local_echo = False
    row = 0
    col = 0
    terminal_type = 'xterm-256color'
    if n > 1 and (args[1] == '-h' or args[1] == '--help'):
        mqtt_out(pty_doc)
        return 0

    if n > 1 and args[1] == '-r':
        raw_input = True
        args = args[1:]
        n -= 1

    if n > 1 and args[1] == '-e':
        local_echo = True
        args = args[1:]
        n -= 1

    if n > 1 and args[1] == '-t':
        terminal_type = args[2]
        args = args[2:]
        n -= 2

    if n > 2:
        row = int(args[1])
        col = int(args[2])
        args = args[2:]

    def hijacker_to_pty(cmd):
        if jc_pty._check_subprocess():
            return True, True
        if raw_input:
            jc_pty.input_into_pty_no_newline(cmd)
        else:
            jc_pty.input_into_pty(cmd)
        return False, False

    def hijacker_from_pty_ext(cmd):
        cmd = cmd.strip()
        if cmd == 'exit':
            if not jc_pty._check_subprocess():
                jc_pty.close_pty()
            lily_hijackers.pop()
            return True
        if cmd == 'reset':
            if not jc_pty._check_subprocess():
                jc_pty.close_pty()
            jc_pty.start_pty(mqtt_blob_str)
            return False
        if cmd == 'echo':
            jc_pty.set_pty_echo(True)
        elif cmd == '-echo':
            jc_pty.set_pty_echo(False)
        elif cmd.startswith('size'):
            args = cmd.split(' ')
            if len(args) == 3:
                jc_pty.set_pty_size(int(args[1]), int(args[2]))
        return False

    lily_add_hijack(hijacker_to_pty)
    lily_add_hijack_ext(hijacker_from_pty_ext)
    jc_pty.start_pty(mqtt_blob_str)
    jc_pty.input_into_pty(f'export TERM={terminal_type}')
    if local_echo:
        jc_pty.input_into_pty(f'stty echo')
    if row and col:
        jc_pty.input_into_pty(f'stty rows {row} cols {col}')

    return 0


def cmd_help(n, args):
    if n == 1:
        out_str = 'lily cmds:\n'
        for c, f in lily_cmds:
            out_str += f'{c}\t'
        mqtt_out(out_str)
        return
    out_str = ''
    if args[1] in lily_docs:
        out_str += f'{args[1]}:\t{lily_docs[args[1]]}\n'

    mqtt_out(out_str)
    return 0


def lily_init():
    new_cmd('help', cmd_help,
            'help [cmd]: show help info, if cmd is given, show help info of cmd')
    new_cmd('ping', lambda n, args: mqtt_out('pong'), 'ping: test connection')
    if sys.platform == 'linux' or sys.platform == 'darwin':
        new_cmd('pty', cmd_pty, pty_doc)
    if not os.path.isdir('lily'):
        return
    for f in os.listdir('lily'):
        if not f.endswith('.py'):
            continue
        try:
            exec(f'import lily.{f[:-3]}')
        except Exception as e:
            print(f'error:{e}')


if __name__ == '__main__':
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    config = configparser.ConfigParser()
    config_file = 'mqtt.ini' if len(sys.argv) <= 1 else sys.argv[1]
    if not os.path.exists(config_file):
        print(
            f'config file {config_file} not found\n usage: {sys.argv[0]} [config_file]')
        exit(-1)
    config.read(config_file)
    client_id = config.get('jc-pi', 'client_id')
    assert client_id
    mqtt_server = config.get(
        'jc-pi', 'server', fallback='aahqtdc.iot.gz.baidubce.com')
    mqtt_username = f'thingidp@aahqtdc|{client_id}'
    mqtt_password = config.get('jc-pi', 'password')
    assert mqtt_password
    mqtt_subscribe_topics = [f'lily/{client_id}/cmd', f'lily/{client_id}/ext', f'lily/{client_id}/ping']
    mqtt_port = config.getint('jc-pi', 'port', fallback=1883)
    mqtt_keep_alive = config.get('jc-pi', 'keep_alive', fallback=120)

    try:
        client = mqtt.Client(client_id)
        client.on_connect = on_connect
        client.on_disconnect = on_disconnect
        client.on_message = on_message
        client.on_connect_fail = lambda a, b, c: exit(-2)
        client.username_pw_set(mqtt_username, mqtt_password)
        client.connect(mqtt_server, mqtt_port, mqtt_keep_alive)
        lily_init()
        for topic in mqtt_subscribe_topics:
            client.subscribe(topic)

        while True:
            try:
                client.loop_forever()
                if not client.is_connected():
                    client.connect(mqtt_server, mqtt_port, mqtt_keep_alive)
                    for topic in mqtt_subscribe_topics:
                        client.subscribe(topic)
                    sleep(2)
                sleep(2)
            except (OSError, mqtt.WebsocketConnectionError) as e:
                print(e)

    except KeyboardInterrupt as e:
        print('stopped')
    except Exception as e:
        print(e)
    finally:
        client.loop_stop()
